{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Searching for Limit Cycles via Trajectory Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import pydrake\n",
    "    import underactuated\n",
    "except ImportError:\n",
    "    !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "    from jupyter_setup import setup_underactuated\n",
    "    setup_underactuated()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# others\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import HTML\n",
    "\n",
    "# drake\n",
    "from pydrake.all import (MultibodyPlant, Parser, DiagramBuilder, Simulator,\n",
    "                         PlanarSceneGraphVisualizer, SceneGraph, TrajectorySource,\n",
    "                         SnoptSolver, MultibodyPositionToGeometryPose, PiecewisePolynomial,\n",
    "                         MathematicalProgram, JacobianWrtVariable, eq)\n",
    "\n",
    "from underactuated import FindResource"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Description\n",
    "\n",
    "In this exercise we will write a nonlinear optimization to find the limit cycle of the passive compass gait.\n",
    "We have already discussed this problem in the [lecture notes](http://underactuated.mit.edu/simple_legs.html#compass_gait) under some simplifying assumptions (the stance foot acts as a pin joint and does not slip, no friction limits in the heel strike, no collision checking between the swing leg and the ground).\n",
    "In the case of the rimless wheel, we have also seen how to use `DirectCollocation` to [detect the limit cycle under these simplifying assumptions](https://github.com/RussTedrake/underactuated/blob/master/examples/limit_cycles.ipynb).\n",
    "In this notebook we take a more general approach, that can be used as a starting point to identify limit cycles for more complex robots, like the [kneed walker](http://underactuated.mit.edu/simple_legs.html#kneed_walker) or even 3D robots.\n",
    "\n",
    "We describe the system in floating-base coordinates, meaning that the position of the stance foot is not fixed and is included in the configuration vector $\\mathbf{q}$.\n",
    "We then parse the [`urdf` file](https://github.com/RussTedrake/underactuated/blob/master/underactuated/models/compass_gait_limit_cycle.urdf) to get the kinematic and dynamic models of the compass gait.\n",
    "Finally, we write a `MathematicalProgram` to identify the limit cycle.\n",
    "\n",
    "This `MathematicalProgram` is a different kind of trajectory optimization from the ones you have been seen before: there is no control input to optimize, the only knob we have is the initial state of the compass gait.\n",
    "We do not even have a cost function to minimize.\n",
    "But, as you will see, this problem is harder than it seems!\n",
    "Your goal is to complete the mathematical program we partially wrote below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Model\n",
    "\n",
    "Let's have a quick look at the model we use.\n",
    "In the figure below we depicted the compass gait with the system of coordinates employed in this exercise.\n",
    "The position of the stance foot, with respect to a frame aligned with the ground, has coordinates $x, y$.\n",
    "The absolute angle of the stance leg is $\\theta_1$, and the angle of the swing leg relative to the stance leg is $\\theta_2$.\n",
    "The configuration vector is $\\mathbf{q} = [x, y, \\theta_1, \\theta_2]^T$; the system state is $\\mathbf{x} = [\\mathbf{q}^T, \\dot{\\mathbf{q}}^T]^T$.\n",
    "The links with mass are the two legs, and the body (white circle).\n",
    "If you are curious to know the numeric parameters of this system, have a look at its [`urdf` file](https://github.com/RussTedrake/underactuated/blob/master/underactuated/models/compass_gait_limit_cycle.urdf).\n",
    "\n",
    "Why this weird system of coordinates and not the one from [lecture notes](http://underactuated.mit.edu/simple_legs.html#compass_gait)?\n",
    "As you will see below, this definition of the configuration vector $\\mathbf{q}$ makes the trajectory optimization easier: many of the constraints we need turn out to be linear in these coordinates.\n",
    "This at the price of some headaches when deriving the periodicity constraints...\n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/RussTedrake/underactuated/master/figures/exercises/compass_gait.jpg\" width=\"500\">\n",
    "\n",
    "The idea is simple: we need to find an initial state so that, after a walking cycle, the robot comes back exactly to the same state.\n",
    "Since the robot is completely symmetric (the legs have the same length and mass distribution), there is no need to optimize for the whole walking cycle (two steps).\n",
    "We just optimize a single step, then we \"mirror\" the result to obtain the second step and complete the walking cycle (see the figure below).\n",
    "\n",
    "<img src=\"https://raw.githubusercontent.com/RussTedrake/underactuated/master/figures/exercises/compass_gait_configurations.jpg\" width=\"500\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Walking Cycle\n",
    "\n",
    "At the initial time $t=0$, we require the robot to start its motion with both the feet on the ground (figure on the left above).\n",
    "\n",
    "The stance foot must be in contact with the ground for all times ($y(t)=0$), while the swing foot is allowed to break contact with the ground but not to penetrate it.\n",
    "The contact force at the stance foot must always lie in the friction cone.\n",
    "\n",
    "To move from the initial to the final configurations above, the compass gait must necessarily scuff the swing foot on the ground at some time between $t=0$ and $t=t_c$.\n",
    "In our model we assume that this scuffing does not generate any contact force.\n",
    "\n",
    "At the end of the step (right figure), the swing foot must collide (i.e. be in contact) with the ground: this moment is called the \"heel strike\" and denoted as $t_c$.\n",
    "At the heel strike we have an impulse on the swing foot which must also lie in the friction cone ($\\int_{t_c^-}^{t_c^+} \\lambda dt$ [in the textbook appendix](http://underactuated.csail.mit.edu/multibody.html#impulsive_collision)).\n",
    "This impulse leads to a jump in the system state as described [in the lecture notes](http://underactuated.mit.edu/simple_legs.html#compass_gait).\n",
    "\n",
    "The periodicity of the motion can be enforced as follows:\n",
    "- Periodicity of the configuration $\\mathbf{q}$.\n",
    "We want the final angles of the legs to be equal to the initial one, but with reversed signs: $\\theta_1(t_c) = \\theta_1(0)$ for the stance leg, and $\\theta_2(t_c) = - \\theta_2(0)$ for the swing leg.\n",
    "- Periodicity of the velocity $\\dot{\\mathbf{q}}$.\n",
    "Since the heel strike is inelastic, after the impact the swing foot must have zero translational velocity.\n",
    "Hence, to enforce periodicity, we require $\\dot{x}(0) = \\dot{y}(0) = 0$.\n",
    "The (absolute) angular velocity of the swing leg at time $t=t_c^+$ must be equal to the one of the stance leg at time $t=0$: some simple kinematics shows that this can be enforced as $\\dot{\\theta}_1(0) = \\dot{\\theta}_1(t_c^+) + \\dot{\\theta}_2(t_c^+)$.\n",
    "Finally, to ensure that the relative velocity of the legs is periodic, we simply enforce $\\dot{\\theta}_2(0) = - \\dot{\\theta}_2(t_c^+)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parse the `urdf` and Get the `MultibodyPlant`\n",
    "\n",
    "We start by defining a couple of physical parameters that we will need below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# friction coefficient between feet and ground\n",
    "friction = .2\n",
    "\n",
    "# position of the feet in the respective leg frame\n",
    "# (must match the urdf)\n",
    "foot_in_leg = {\n",
    "    'stance_leg': np.zeros(3),        # stance foot in stance-leg frame\n",
    "    'swing_leg': np.array([0, 0, -1]) # swing foot in swing-leg frame\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we parse the `urdf` file to get a Drake `MultibodyPlant`.\n",
    "\n",
    "**The swing-foot scuffing.**\n",
    "In this notebook we will use the compass gait `MultibodyPlant` only as a support for the computation of the robot kinematics and dynamics.\n",
    "In particular, we will not use it for simulation.\n",
    "The reason for this resides in the scuffing issue described above: when the swing foot is sliding on the ground at zero height, any good physic simulator detects a collision and applies a friction force to the foot.\n",
    "This would cause the compass gait ot stumble and fall down, no matter the trajectory we found.\n",
    "Note however that more complex robots, such as the [kneed walker](http://underactuated.mit.edu/simple_legs.html#kneed_walker), do not have this issue, and, for these, the workflow we use in this notebook can be used both for the detection of the limit cycle and for simulation.\n",
    "\n",
    "**An important implementation detail.**\n",
    "Behind the scenes, the optimization solvers we use require the knowledge of the derivaties of the cost function and the constraint values with respect to the decision variables.\n",
    "This is needed to understand in which direction the current solution should be corrected to find a feasible point or reduce the cost.\n",
    "This process used to be very tedious some years ago, when graduate students had to spend many hours writing down these derivatives \"by hand\".\n",
    "Nowadays, we use [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation), which through the construction of a computational graph is able to evaluate a function and its derivatives very quickly and exactly (cf. [finite difference](https://en.wikipedia.org/wiki/Finite_difference)).\n",
    "To allow the evaluation of the `MultibodyPlant` functions (e.g. the mass matrix method) with `AutoDiffXd` variables, we need to call the `MultibodyPlant.ToAutoDiffXd()` function which returns a copy of the `MultibodyPlant` that can work with autodiff variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parse urdf and create the MultibodyPlant\n",
    "compass_gait = MultibodyPlant(time_step=0)\n",
    "file_name = FindResource('models/compass_gait_limit_cycle.urdf')\n",
    "Parser(compass_gait).AddModelFromFile(file_name)\n",
    "compass_gait.Finalize()\n",
    "\n",
    "# overwrite MultibodyPlant with its autodiff copy\n",
    "compass_gait = compass_gait.ToAutoDiffXd()\n",
    "\n",
    "# number of configuration variables\n",
    "nq = compass_gait.num_positions()\n",
    "\n",
    "# number of components of the contact forces\n",
    "nf = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helper Functions for the `MathematicalProgram`\n",
    "\n",
    "When writing a `MathematicalProgram` in Drake, optimization variables are `symbolic.Variable` objects.\n",
    "These cannot be passed directly to the `MultibodyPlant` functions (such as `CalcMassMatrix`), which only accept floats or `AutoDiffXd` types.\n",
    "Hence, if you need to add a constraint which involves the evaluation of a `MultibodyPlant` function, you need to proceed as follows.\n",
    "\n",
    "Write a python function (say `my_fun`) that, given the numeric value (`float` or `AutoDiffXd`) of certain variables in the problem, returns the numeric value of the quantity that you want to constrain.\n",
    "Let `vars` be the arguments of this function and `values` its output (both can be arrays).\n",
    "Using the method `MathematicalProgram.AddConstraint` you can write `prog.AddConstraint(my_fun, lb=values_lb, ub=values_ub, vars=vars)` to enforce the constraints `values_lb <= values <= values_ub`, where `value_lb` and `value_ub` are vectors of floats of appropriate dimensions.\n",
    "Then, at solution time, the solver will evaluate `my_fun` passing autodiff variables, retrieving in this way the numeric values of the constraint violations and their derivatives.\n",
    "\n",
    "In the following cell we wrote the functions that we will need to enforce the necessary constraints in the trajectory optimization problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function that given the current configuration, velocity,\n",
    "# acceleration, and contact force at the stance foot, evaluates\n",
    "# the manipulator equations. The output of this function is a\n",
    "# vector with dimensions equal to the number of configuration\n",
    "# variables. If the output of this function is equal to zero\n",
    "# then the given arguments verify the manipulator equations.\n",
    "def manipulator_equations(vars):\n",
    "    \n",
    "    # split input vector in subvariables\n",
    "    # configuration, velocity, acceleration, stance-foot force\n",
    "    assert vars.size == 3 * nq + nf\n",
    "    split_at = [nq, 2 * nq, 3 * nq]\n",
    "    q, qd, qdd, f = np.split(vars, split_at)\n",
    "    \n",
    "    # set compass gait state\n",
    "    context = compass_gait.CreateDefaultContext()\n",
    "    compass_gait.SetPositions(context, q)\n",
    "    compass_gait.SetVelocities(context, qd)\n",
    "    \n",
    "    # matrices for the manipulator equations\n",
    "    M = compass_gait.CalcMassMatrixViaInverseDynamics(context)\n",
    "    Cv = compass_gait.CalcBiasTerm(context)\n",
    "    tauG = compass_gait.CalcGravityGeneralizedForces(context)\n",
    "    \n",
    "    # Jacobian of the stance foot\n",
    "    J = get_foot_jacobian(compass_gait, context, 'stance_leg')\n",
    "    \n",
    "    # return violation of the manipulator equations\n",
    "    return M.dot(qdd) + Cv - tauG - J.T.dot(f)\n",
    "\n",
    "# Function that given the current configuration, returns\n",
    "# the distance of the swing foot from the ground (scalar).\n",
    "# We have penetration if the function output is negative.\n",
    "def swing_foot_height(q):\n",
    "    \n",
    "    # get reference frames for the swing leg and the ground\n",
    "    leg_frame = compass_gait.GetBodyByName('swing_leg').body_frame()\n",
    "    ground_frame = compass_gait.GetBodyByName('ground').body_frame()\n",
    "    \n",
    "    # position of the swing foot in ground coordinates\n",
    "    context = compass_gait.CreateDefaultContext()\n",
    "    compass_gait.SetPositions(context, q)\n",
    "    swing_foot_position = compass_gait.CalcPointsPositions(\n",
    "        context,\n",
    "        leg_frame,\n",
    "        foot_in_leg['swing_leg'],\n",
    "        ground_frame\n",
    "    )\n",
    "    \n",
    "    # return only the coordinate z\n",
    "    # (distance normal to the ground)\n",
    "    return swing_foot_position[-1]\n",
    "\n",
    "# Function that implements the impulsive collision derived in\n",
    "# the textbook appendix. Arguments are: compass gait configuration,\n",
    "# velocities before and after heel strike, and the swing-foot\n",
    "# impulse (in latex, $\\int_{t_c^-}^{t_c^+} \\lambda dt$).\n",
    "# Returns a vector of quantities that must vanish in order\n",
    "# for the impulsive dynamics to be verified: it enforces the velocity\n",
    "# jump due to the impulse, and the inelastic behavior of the\n",
    "# collision (zero coefficient of restitution $e$).\n",
    "# See http://underactuated.mit.edu/multibody.html#impulsive_collision\n",
    "def reset_velocity_heelstrike(vars):\n",
    "    \n",
    "    # split input vector in subvariables\n",
    "    # qd_pre: generalized velocity before the heel strike\n",
    "    # qd_post: generalized velocity after the heel strike\n",
    "    # imp: swing-foot collision impulse (2d vector)\n",
    "    assert vars.size == 3 * nq + nf\n",
    "    split_at = [nq, 2 * nq, 3 * nq]\n",
    "    q, qd_pre, qd_post, imp = np.split(vars, split_at)\n",
    "\n",
    "    # set compass gait configuration\n",
    "    context = compass_gait.CreateDefaultContext()\n",
    "    compass_gait.SetPositions(context, q)\n",
    "    \n",
    "    # get necessary matrices\n",
    "    M = compass_gait.CalcMassMatrixViaInverseDynamics(context)\n",
    "    J = get_foot_jacobian(compass_gait, context, 'swing_leg')\n",
    "    \n",
    "    # return a vector that must vanish for the impulsive dynamics to hold\n",
    "    return np.concatenate((\n",
    "        M.dot(qd_post - qd_pre) - J.T.dot(imp), # velocity jump due to the impulse\n",
    "        J.dot(qd_post)                          # zero velocity restitution\n",
    "    ))\n",
    "\n",
    "# Function that given a leg, returns the Jacobian matrix for the related foot.\n",
    "def get_foot_jacobian(compass_gait, context, leg):\n",
    "    \n",
    "    # get reference frames for the given leg and the ground\n",
    "    leg_frame = compass_gait.GetBodyByName(leg).body_frame()\n",
    "    ground_frame = compass_gait.GetBodyByName('ground').body_frame()\n",
    "\n",
    "    # compute Jacobian matrix\n",
    "    J = compass_gait.CalcJacobianTranslationalVelocity(\n",
    "        context,\n",
    "        JacobianWrtVariable(0),\n",
    "        leg_frame,\n",
    "        foot_in_leg[leg],\n",
    "        ground_frame,\n",
    "        ground_frame\n",
    "    )\n",
    "    \n",
    "    # discard y components since we are in 2D\n",
    "    return J[[0, 2]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Trajectory Optimization Problem\n",
    "\n",
    "We start by setting some parameters of our optimization problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# time steps in the trajectory optimization\n",
    "T = 50\n",
    "\n",
    "# minimum and maximum time interval is seconds\n",
    "h_min = .005\n",
    "h_max = .05"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Troubleshooting.**\n",
    "To simplify the reading, we divide the construction of the `MathematicalProgram` in multiple cells.\n",
    "If you modify any of the components of the problem, be sure to rerun your code starting from the following cell (where the `MathematicalProgram` is initialized).\n",
    "Otherwise you will add the same constraints multiple times to the same optimization problem.\n",
    "\n",
    "We start from the decision variables of the trajectory optimization problem.\n",
    "Notice that we also add the accelerations `qdd` among the optimization variables here.\n",
    "This is slightly unusual, and not necessary, but in these circumstances it simplifies a bit the code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize program\n",
    "prog = MathematicalProgram()\n",
    "\n",
    "# vector of the time intervals\n",
    "# (distances between the T + 1 break points)\n",
    "h = prog.NewContinuousVariables(T, name='h')\n",
    "\n",
    "# system configuration, generalized velocities, and accelerations\n",
    "q = prog.NewContinuousVariables(rows=T+1, cols=nq, name='q')\n",
    "qd = prog.NewContinuousVariables(rows=T+1, cols=nq, name='qd')\n",
    "qdd = prog.NewContinuousVariables(rows=T, cols=nq, name='qdd')\n",
    "\n",
    "# stance-foot force\n",
    "f = prog.NewContinuousVariables(rows=T, cols=nf, name='f')\n",
    "\n",
    "# heel strike impulse for the swing leg\n",
    "imp = prog.NewContinuousVariables(nf, name='imp')\n",
    "\n",
    "# generalized velocity after the heel strike\n",
    "# (if \"mirrored\", this velocity must coincide with the\n",
    "# initial velocity qd[0] to ensure periodicity)\n",
    "qd_post = prog.NewContinuousVariables(nq, name='qd_post')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are part of the constraints of the optimization problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# lower and upper bound on the time steps for all t\n",
    "prog.AddBoundingBoxConstraint([h_min] * T, [h_max] * T, h)\n",
    "\n",
    "# link the configurations, velocities, and accelerations\n",
    "# uses implicit Euler method, https://en.wikipedia.org/wiki/Backward_Euler_method\n",
    "for t in range(T):\n",
    "    prog.AddConstraint(eq(q[t+1], q[t] + h[t] * qd[t+1]))\n",
    "    prog.AddConstraint(eq(qd[t+1], qd[t] + h[t] * qdd[t]))\n",
    "\n",
    "# manipulator equations for all t (implicit Euler)\n",
    "for t in range(T):\n",
    "    vars = np.concatenate((q[t+1], qd[t+1], qdd[t], f[t]))\n",
    "    prog.AddConstraint(manipulator_equations, lb=[0]*nq, ub=[0]*nq, vars=vars)\n",
    "    \n",
    "# velocity reset across heel strike\n",
    "# see http://underactuated.mit.edu/multibody.html#impulsive_collision\n",
    "vars = np.concatenate((q[-1], qd[-1], qd_post, imp))\n",
    "prog.AddConstraint(reset_velocity_heelstrike, lb=[0]*(nq+nf), ub=[0]*(nq+nf), vars=vars)\n",
    "    \n",
    "# mirror initial and final configuration\n",
    "# see \"The Walking Cycle\" section of this notebook\n",
    "prog.AddLinearConstraint(eq(q[0], - q[-1]))\n",
    "\n",
    "# mirror constraint between initial and final velocity\n",
    "# see \"The Walking Cycle\" section of this notebook\n",
    "prog.AddLinearConstraint(qd[0, 0] == 0)\n",
    "prog.AddLinearConstraint(qd[0, 1] == 0)\n",
    "prog.AddLinearConstraint(qd[0, 2] == qd_post[2] + qd_post[3])\n",
    "prog.AddLinearConstraint(qd[0, 3] == - qd_post[3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it is your turn to complete the optimization problem.\n",
    "You need to add five groups of constraints:\n",
    "1. **Stance foot on the ground for all times.**\n",
    "This `LinearConstraint` must ensure that $x(t) = y(t) = 0$ for all $t$.\n",
    "2. **Swing foot on the ground at time zero.**\n",
    "This constraint must ensure that the initial configuration $\\mathbf{q}(0)$ is such that the swing foot is on the ground.\n",
    "For a more complex robot, this would be a tough nonlinear constraint.\n",
    "For the compass gait, knowing that the stance foot is on the ground, you should be able to express this as a `LinearConstraint` on the entries of $\\mathbf{q}(0)$.\n",
    "3. **No penetration of the swing foot in the ground for all times.**\n",
    "This nonlinear constraint can be added using the function `swing_foot_height` we defined above.\n",
    "Follow the examples from the previous cell to see how to add a nonlinear constraint defined via a python function using `AddConstraint`.\n",
    "Note that, in this case, you want the upper bound on the function output to be `ub=[np.inf]` for all times.\n",
    "4. **Stance-foot contact force in friction cone for all times.**\n",
    "To prevent the robot from slipping, the contact force `f[t]` must lie in the friction cone for all `t`.\n",
    "This means that the normal component `f[t, 1]` must be nonnegative, and the tangential component `f[t, 0]` must be, in absolute value, lower or equal than the normal component `f[t, 1]` times the friction coefficient (`friction` here).\n",
    "Note that these conditions can be enforced as three `LinearConstraint`s per time step `t`.\n",
    "5. **Swing-foot impulse in friction cone.**\n",
    "To ensure that the swing foot is completely stopped by the heel strike, ensure that the impulse `imp` belongs to the friction cone.\n",
    "This can be done using `AddLinearConstraint` three times, as for the previous bullet point.\n",
    "\n",
    "**Troubleshooting:**\n",
    "Unfortunately, nonlinear solvers are very sensitive.\n",
    "It can happen that, just changing the order in which you add constraints to the problem, you get a different solution.\n",
    "We suggest you to add these constraints in the given order: for us it worked fine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. stance foot on the ground for all times\n",
    "# modify here\n",
    "\n",
    "# 2. swing foot on the ground at time zero\n",
    "# modify here\n",
    "\n",
    "# 3. no penetration of the swing foot in the ground for all times\n",
    "# modify here\n",
    "\n",
    "# 4. stance-foot contact force in friction cone for all times\n",
    "# modify here\n",
    "\n",
    "# 5. swing-foot impulse in friction cone\n",
    "# modify here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we set the initial guess for our optimization problem.\n",
    "\n",
    "For the time steps `h` we just initialize them to their maximal value `h_max` (somewhat an arbitrary decision, but it works).\n",
    "\n",
    "For the robot configuration `q`, we interpolate between the initial value `q0_guess` and the final value `- q0_guess`.\n",
    "In our implementation, the value given below for `q0_guess` made the optimization converge.\n",
    "But, if you find the need, feel free to tweak this parameter.\n",
    "The initial guess for the velocity and the acceleration is obtained by differentiating the one for the position.\n",
    "\n",
    "The normal force `f` at the stance foot is equal to the total `weight` of the robot.\n",
    "\n",
    "All the other optimization variables are initialized at zero.\n",
    "(Note that, if the initial guess for a variable is not specified, the default value is zero.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# vector of the initial guess\n",
    "initial_guess = np.empty(prog.num_vars())\n",
    "\n",
    "# initial guess for the time step\n",
    "h_guess = h_max\n",
    "prog.SetDecisionVariableValueInVector(h, [h_guess] * T, initial_guess)\n",
    "\n",
    "# linear interpolation of the configuration\n",
    "q0_guess = np.array([0, 0, .15, -.3])\n",
    "q_guess_poly = PiecewisePolynomial.FirstOrderHold(\n",
    "    [0, T * h_guess],\n",
    "    np.column_stack((q0_guess, - q0_guess))\n",
    ")\n",
    "qd_guess_poly = q_guess_poly.derivative()\n",
    "qdd_guess_poly = q_guess_poly.derivative()\n",
    "\n",
    "# set initial guess for configuration, velocity, and acceleration\n",
    "q_guess = np.hstack([q_guess_poly.value(t * h_guess) for t in range(T + 1)]).T\n",
    "qd_guess = np.hstack([qd_guess_poly.value(t * h_guess) for t in range(T + 1)]).T\n",
    "qdd_guess = np.hstack([qdd_guess_poly.value(t * h_guess) for t in range(T)]).T\n",
    "prog.SetDecisionVariableValueInVector(q, q_guess, initial_guess)\n",
    "prog.SetDecisionVariableValueInVector(qd, qd_guess, initial_guess)\n",
    "prog.SetDecisionVariableValueInVector(qdd, qdd_guess, initial_guess)\n",
    "\n",
    "# initial guess for the normal component of the stance-leg force\n",
    "bodies = ['body', 'stance_leg', 'swing_leg']\n",
    "mass = sum(compass_gait.GetBodyByName(body).default_mass() for body in bodies)\n",
    "g = - compass_gait.gravity_field().gravity_vector()[-1]\n",
    "weight = mass * g\n",
    "prog.SetDecisionVariableValueInVector(f[:, 1], [weight] * T, initial_guess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can finally solve the problem! Be sure that the solver actually converged: you can check this by looking at the variable `result.is_success()` (printed below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# solve mathematical program with initial guess\n",
    "solver = SnoptSolver()\n",
    "result = solver.Solve(prog, initial_guess)\n",
    "\n",
    "# ensure solution is found\n",
    "print(f'Solution found? {result.is_success()}.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following cell we retrieve the optimal value of the decision variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get optimal solution\n",
    "h_opt = result.GetSolution(h)\n",
    "q_opt = result.GetSolution(q)\n",
    "qd_opt = result.GetSolution(qd)\n",
    "qdd_opt = result.GetSolution(qdd)\n",
    "f_opt = result.GetSolution(f)\n",
    "imp_opt = result.GetSolution(imp)\n",
    "qd_post_opt = result.GetSolution(qd_post)\n",
    "\n",
    "# stack states\n",
    "x_opt = np.hstack((q_opt, qd_opt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Animate the Result\n",
    "\n",
    "Here we quickly build a Drake diagram to animate the result we got from trajectory optimization: useful for debugging your code and to be sure that everything looks good."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# interpolate state values for animation\n",
    "time_breaks_opt = np.array([sum(h_opt[:t]) for t in range(T+1)])\n",
    "x_opt_poly = PiecewisePolynomial.FirstOrderHold(time_breaks_opt, x_opt.T)\n",
    "\n",
    "# parse urdf with scene graph\n",
    "compass_gait = MultibodyPlant(time_step=0)\n",
    "scene_graph = SceneGraph()\n",
    "compass_gait.RegisterAsSourceForSceneGraph(scene_graph)\n",
    "file_name = FindResource('models/compass_gait_limit_cycle.urdf')\n",
    "Parser(compass_gait).AddModelFromFile(file_name)\n",
    "compass_gait.Finalize()\n",
    "\n",
    "# build block diagram and drive system state with\n",
    "# the trajectory from the optimization problem\n",
    "builder = DiagramBuilder()\n",
    "source = builder.AddSystem(TrajectorySource(x_opt_poly))\n",
    "builder.AddSystem(scene_graph)\n",
    "pos_to_pose = builder.AddSystem(MultibodyPositionToGeometryPose(compass_gait, input_multibody_state=True))\n",
    "builder.Connect(source.get_output_port(0), pos_to_pose.get_input_port())\n",
    "builder.Connect(pos_to_pose.get_output_port(), scene_graph.get_source_pose_port(compass_gait.get_source_id()))\n",
    "\n",
    "# add visualizer\n",
    "xlim = [-.75, 1.]\n",
    "ylim = [-.2, 1.5]\n",
    "visualizer = builder.AddSystem(PlanarSceneGraphVisualizer(scene_graph, xlim=xlim, ylim=ylim, show=False))\n",
    "builder.Connect(scene_graph.get_pose_bundle_output_port(), visualizer.get_input_port(0))\n",
    "simulator = Simulator(builder.Build())\n",
    "\n",
    "# generate and display animation\n",
    "visualizer.start_recording()\n",
    "simulator.AdvanceTo(x_opt_poly.end_time())\n",
    "ani = visualizer.get_recording_as_animation()\n",
    "HTML(ani.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the Results\n",
    "\n",
    "Here are two plots to visualize the results of the trajectory optimization.\n",
    "\n",
    "In the first we plot the limit cycle we found in the plane of the leg angles.\n",
    "To show a complete cycle, we \"mirror\" the trajectory of the first step and we plot it too (\"Red leg swinging\")."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot swing trajectories\n",
    "# the second is the mirrored one\n",
    "plt.plot(q_opt[:, 2], q_opt[:, 3], color='b', label='Blue leg swinging')\n",
    "plt.plot(q_opt[:, 2] + q_opt[:, 3], - q_opt[:, 3], color='r', label='Red leg swinging')\n",
    "\n",
    "# scatter heel strikes\n",
    "plt.scatter(q_opt[0, 2] + q_opt[0, 3], - q_opt[0, 3], color='b', zorder=3, label='Blue-leg heel strike')\n",
    "plt.scatter(q_opt[0, 2], q_opt[0, 3], color='r', zorder=3, label='Red-leg heel strike')\n",
    "\n",
    "# misc options\n",
    "plt.xlabel('Absolute angle red leg')\n",
    "plt.ylabel('Angle blue leg wrt red leg')\n",
    "plt.grid(True)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we plot the absolute angle of the red leg versus its time derivative.\n",
    "Again we complete the cycle by mirroring the result of the trajectory step.\n",
    "\n",
    "If you did thing correctly, this figure should resemble [Figure 4.10 from the lecture notes](http://underactuated.mit.edu/simple_legs.html#compass_gait) (with reversed signs).\n",
    "Note that the angle of the red leg is continuous during the walking cycle, while its time derivative has two jumps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot swing trajectories\n",
    "# the second is the mirrored one\n",
    "plt.plot(q_opt[:, 2], qd_opt[:, 2], color='b', label='Blue leg swinging')\n",
    "plt.plot(q_opt[:, 2] + q_opt[:, 3], qd_opt[:, 2] + qd_opt[:, 3], color='r', label='Red leg swinging')\n",
    "\n",
    "# plot heel strikes\n",
    "plt.plot(\n",
    "    [q_opt[-1, 2], q_opt[0, 2] + q_opt[0, 3]],\n",
    "    [qd_opt[-1, 2], qd_opt[0, 2] + qd_opt[0, 3]],\n",
    "    linestyle=':',\n",
    "    color='b',\n",
    "    label='Blue-leg heel strike'\n",
    ")\n",
    "plt.plot(\n",
    "    [q_opt[0, 2], q_opt[-1, 2] + q_opt[-1, 3]],\n",
    "    [qd_opt[0, 2], qd_opt[-1, 2] + qd_opt[-1, 3]],\n",
    "    linestyle=':',\n",
    "    color='r',\n",
    "    label='Red-leg heel strike'\n",
    ")\n",
    "\n",
    "# misc options\n",
    "plt.xlabel('Absolute angle red leg')\n",
    "plt.ylabel('Absolute angular velocity red leg')\n",
    "plt.grid(True)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autograding\n",
    "You can check your work by running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from underactuated.exercises.simple_legs.compass_gait_limit_cycle.test_compass_gait_limit_cycle import TestCompassGaitLimitCycle\n",
    "from underactuated.exercises.grader import Grader\n",
    "Grader.grade_output([TestCompassGaitLimitCycle], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}